//go:build windows

package windows

/*
#cgo CFLAGS: -DINITGUID -w
#cgo LDFLAGS: -ld3d11 -ldxgi -ld3dcompiler -lole32

#include "DirectX11.h"
*/
import "C"
import (
	"errors"
	"image"
	"image/color"
	"sync"
	"syscall"
	"unsafe"
)

type Monitor struct {
	adapter uint
	output  uint
	name    string
	left    int
	right   int
	top     int
	bottom  int
	width   uint
	height  uint
}

type DirectX11Session struct {
	session  *C.DirectX11Session
	monitors []*Monitor

	display uint
	rgb     bool
}

func NewDirectX11Session(display uint, rgb bool) (session *DirectX11Session, err error) {
	monitors, err := GetMonitors()
	if err != nil {
		return
	}
	session = &DirectX11Session{
		session:  &C.DirectX11Session{},
		monitors: monitors,
		display:  display,
		rgb:      rgb,
	}
	result := C.NewDirectX11Session(session.session, C.UINT(display), C.bool(rgb))
	if result != C.S_OK {
		return nil, HResultError(result)
	}

	return
}

func (s *DirectX11Session) Close() {
	C.CloseDirectX11Session(s.session)
	s.display = 0
}

func (s *DirectX11Session) Capture() (img *image.RGBA, err error) {
	var cBits *C.BYTE
	var x, y = 0, 0

	width := s.monitors[s.display].width
	height := s.monitors[s.display].height

	result := C.Screenshot(s.session, C.INT(x), C.INT(y), (*C.UINT)(unsafe.Pointer(&width)), (*C.UINT)(unsafe.Pointer(&height)), &cBits)
	if result != C.S_OK {
		return nil, HResultError(result)
	}

	ptr := (*[1 << 30]byte)(unsafe.Pointer(cBits))[:width*height*4]
	img = image.NewRGBA(image.Rect(0, 0, int(width), int(height)))
	copy(img.Pix, ptr)

	// 释放内存
	C.free(unsafe.Pointer(cBits))

	return
}

func (s *DirectX11Session) switchBGRAtoRGBA(img *image.RGBA) {
	bounds := img.Bounds()
	width := bounds.Dx()
	height := bounds.Dy()

	base := 10
	xAvg := width / base
	yAvg := height / base
	xRange := make([][]int, base)
	yRange := make([][]int, base)
	for n := 0; n < base; n++ {
		start := xAvg * n
		end := xAvg * (n + 1)
		if width-end < xAvg {
			end = width
		}
		xRange[n] = []int{start, end}
	}
	for n := 0; n < base; n++ {
		start := yAvg * n
		end := yAvg * (n + 1)
		if width-end < yAvg {
			end = height
		}
		yRange[n] = []int{start, end}
	}

	wg := &sync.WaitGroup{}
	for _, yIndex := range yRange {
		for _, xIndex := range xRange {
			wg.Add(1)
			go func(xIndex, yIndex []int) {
				for y := yIndex[0]; y < yIndex[1]; y++ {
					for x := xIndex[0]; x < xIndex[1]; x++ {
						r, g, b, a := img.At(x, y).RGBA()
						// 转换为RGBA（交换R和B）
						img.Set(x, y, color.RGBA{
							R: uint8(b >> 8),
							G: uint8(g >> 8),
							B: uint8(r >> 8),
							A: uint8(a >> 8),
						})
					}
				}

				wg.Done()
			}(xIndex, yIndex)
		}
	}
	wg.Wait()
}

func GetScreenRect(display uint) (left, right, top, bottom int, err error) {
	monitors, err := GetMonitors()
	if err != nil {
		return
	}

	return monitors[display].left, monitors[display].right, monitors[display].top, monitors[display].bottom, nil
}

func GetMonitors() (monitors []*Monitor, err error) {
	var adapters C.UINT
	var hr C.HRESULT

	hr = C.CountAdapters(&adapters)
	if !errors.Is(HResultError(hr), S_OK) && !errors.Is(HResultError(hr), DXGI_ERROR_NOT_FOUND) {
		return nil, HResultError(hr)
	}

	for adapter := C.UINT(0); adapter < adapters; adapter++ {
		var output C.UINT
		for output = 0; ; output++ {
			var left, right, top, bottom C.INT
			var buf = make([]uint16, 256)
			hr = C.EnumOutputs(
				adapter,
				output,
				&left,
				&right,
				&top,
				&bottom,
				(*C.WCHAR)(unsafe.Pointer(&buf[0])),
				C.UINT(len(buf)),
			)
			if errors.Is(HResultError(hr), DXGI_ERROR_NOT_FOUND) {
				break
			}
			if !errors.Is(HResultError(hr), S_OK) {
				continue
			}
			monitor := &Monitor{
				adapter: uint(adapter),
				output:  uint(output),
				name:    syscall.UTF16ToString(buf),
				left:    int(left),
				right:   int(right),
				top:     int(top),
				bottom:  int(bottom),
				width:   uint(right - left),
				height:  uint(bottom - top),
			}
			monitors = append(monitors, monitor)
		}
	}

	return
}
